import sys,numpy as np;
# sys.path.append("/home/alfredo/Desktop/libdai-pythonswig/swig/build/lib.linux-x86_64-2.7");
# sys.path.append("/home/alfredo/Desktop/libDAI-0.3.1/swig");
# import dwrap, dai;
# from utils import NodeNames as NN;
from decimal import Decimal, ROUND_DOWN;
import math;

class Node:
    def __init__(self, name="", nValues=2, sId=0, actName="", aSubjects=[], aActorsInfo=[], 
				location="", timeId=0, aTimes=[], nRealInitT=0, aScores=[], sParentsId=[]):
        self.name = name;
        self.nValues = nValues;
        self.sId = sId;
        self.actName = actName;
        self.aSubjects = aSubjects;
        self.aOnlySubjects = [];
        for subject in aSubjects:
            if subject=="*": continue;
            self.aOnlySubjects.append(subject);
        self.aActorsAll = aActorsInfo[0] if len(aActorsInfo)>0 else [];
        self.aActorTypes = aActorsInfo[1] if len(aActorsInfo)>1 else [];
        self.aActorTimes = aActorsInfo[2] if len(aActorsInfo)>2 else [];
        self.aActorRoles = aActorsInfo[3] if len(aActorsInfo)>3 else [];
        self.location = location;
        self.aTimes = aTimes;
        self.nRealInitT = nRealInitT;
        self.aScores = aScores;
        self.sParentsId = sParentsId;
#         self.nForward = 0; #Forward = Likelihood * max[Trans * PastFwd]
        self.mForward = dict(); #Forward = Likelihood * max[Trans * PastFwd]. (one for each HMM of each subject that appears in this node)
        self.mMaxPast_NodeId = dict(); #NodeID to Max Past (Parent) Node (one for each HMM of each subject that appears in this node)
        self.mChildren = dict(); #Children Nodes on Consistency Choice Tree (one for each tuple of HMM's Subject & Parent)
#         self.mJntFwd

class Graph:
    def __init__(self):
        self.ndict = dict()
        self.idndict = dict();
        self.obs = []
            
    def addnode(self,node,parents,cpt):
        self.idndict[node.sId] = node;
        self.ndict[node] = [parents, cpt]
        return(0)
            
    def removeNode(self, node):
        if self.idndict.has_key(node.sId):
            del self.idndict[node.sId];
        if self.ndict.has_key(node):
            del self.ndict[node];
        return(0);
 
    def addobs(self, node, value=False, prior=0.5):
        for i in range(len(self.obs)):
            if self.obs[i][0] == node:
                self.obs[i] = [node, value, prior];
                return(0);
        self.obs.append([node, value, prior]);
        return(0)
    
    def removeObs(self, node): #, value=False, prior=0.5):
        for i in range(len(self.obs)):
            if self.obs[i][0] == node:
                self.obs.remove(self.obs[i]);
                return(0);
            
    def myp2cpt(self, factor, aParents, aCPT, mNode2label):
        def getIndex(aVals, aBase, aIndex):
                nIndex = 0;
                for i, nV in enumerate(aVals):
                    nF, nBases = 1, len(aBase[i+1:]);
                    for j in range(nBases):
                        nF *= aBase[aIndex[j+1]];
                    nIndex += (nV*nF);
                return nIndex;            
        def getPermutation(nIndex, aBase):
            aVals = [];
            for nB in aBase[::-1]:
                aVals.append(nIndex % nB);
                nIndex /= nB;
            aVals.reverse();
            return aVals;
        
        aV = [2,2,2];
        nVals = 1;
        for n in aV:
            nVals *= n;
#         aFrom, aTo = [0,1,2], [2,0,1];
#         aDiff = [aTo.index(i) for i in range(len(aTo))]; #[1,2,0]
#         aV2 = [aV[i] for i in aTo]; #[2,3,2]
        ids, aTo = [], [];
        for parent in aParents:
            ids.append(mNode2label[parent])
            aTo.append(mNode2label[parent])
        aCPT = permute(self, aCPT,ids)
        print "aTo:", aTo;
        
        plen = len(aCPT)
#         tset = [0]*2*plen
        tlen = np.log2(plen)+1 ## length our binary string
        for i,entry in enumerate(aCPT):
            entry = 1-entry
            tstring = bin(i) ## this string has format '0bxxxxx'
            tstring = tstring[2:] ## strip off leading "0b"
            tstring = "0"*(tlen-len(tstring))+tstring
            tstring = tstring[1:] ## don't want to reverse first digit
            rstring = tstring[::-1]
            ## reverse the slice, this gives us the binary number we want
            s1 = "0"+rstring
            ind = int(s1,2) ## where in the list we want to put this prob
            factor[ind] = entry
            factor[ind+plen] = 1-entry
        return factor

    def mySetup(self):
        #Create Nodes for libDAI
        for nId,node in enumerate(self.ndict):
            self.node2label[node] = nId;
            self.node2daiVar[node] = dai.Var(nId, node.nValues);
        #Create Factors with these Nodes & their Parents
        vecFactors = dai.VecFactor();
        for node in self.node2daiVar:
            varSet = dai.VarSet(self.node2daiVar[node]);
            for parent in node[0]:
                varSet.append(parent);
            factor = dai.Factor(varSet);
            parents, aCPT = node[0], node[1];
            vecFactors.append(self.myp2cpt(factor, parents, aCPT, self.node2label));
        self.daiNetwork = dai.FactorGraph(vecFactors);
        self.daiNetwork.WriteToFile("tempdai.fg");
        
    ##These are suboptimal, pls don't use with too many nodes
    def setup(self):
        self.setup2();
        return;
        def toposort(nodedict):
            visited = dict()
            nodelist = []
            def visit(node):
#                print node.name;
                try:
                    visited[node]
#                    print "visited!";
                except KeyError:
                    visited[node] = 1                    
                    pset = nodedict[node][0] ## parents                    
                    for p in pset:
                        visit(p)
#                    print "Appending:", node.name;
                    nodelist.append(node)
            for node in nodedict.keys():
#                print "Starting with:", node.name;
                visit(node)
#            print "";
#            for node in nodelist:
#                print "Ending with:", node.name;
            return(nodelist)
        def p2cpt(pset, parents, ptable, node2label):
            def getIndex(aVals, aBase, aIndex):
                nIndex = 0;
                for i, nV in enumerate(aVals):
                    nF, nBases = 1, len(aBase[i+1:]);
                    for j in range(nBases):
                        nF *= aBase[aIndex[j+1]];
                    nIndex += (nV*nF);
                return nIndex;            
            def getPermutation(nIndex, aBase):
                aVals = [];
                for nB in aBase[::-1]:
                    aVals.append(nIndex % nB);
                    nIndex /= nB;
                aVals.reverse();
                return aVals;
            def permute(self,ptable, ids):
                order = [i[0] for i in sorted(enumerate(ids), key=lambda x:x[1])]
                pnew = [0]*len(ptable)
                for (i,val) in enumerate(ptable):    
                    d = self.tobin(i, len(ids))
                    od = [d[j] for j in order]
                    newi = self.frombin(od)
                    if newi>len(pnew)-1:
#                        newi = len(pnew)-1;
                        print "ERROR (newi>len(pnew)-1): Node's CPT has incorrect number of probability arguments.";
                    pnew[newi] = val
                return pnew            
            aV = [2,2,2];
            nVals = 1;
            for n in aV:
                nVals *= n;
#             aFrom, aTo = [0,1,2], [2,0,1];
#             aDiff = [aTo.index(i) for i in range(len(aTo))]; #[1,2,0]
#             aV2 = [aV[i] for i in aTo]; #[2,3,2]

            ids, aTo = [], [];
            for parent in parents:
                ids.append(node2label[parent])
                aTo.append(node2label[parent])
            ptable = permute(self, ptable,ids)
#            print "aTo:", aTo;

            plen = len(ptable)
#             tset = [0]*2*plen
            tlen = np.log2(plen)+1 ## length our binary string
            for i,entry in enumerate(ptable):
                entry = 1-entry
                ## entries all assume entry conditioned on others is 0
                ## ptables count up across the remaining variables in binary
                ## so e.g. for 3 variables (c|a,b), the table looks like
                ## 0 0 0
                ## 0 0 1
                ## 0 1 0
                ## 0 1 1
                ## however, dai requires that the full table be specified, and
                ## counts in "reversed" binary, so the full table for the above
                ## example would be ordered:
                ## 0 0 0
                ## 1 0 0
                ## 0 1 0
                ## 1 1 0
                ## 0 0 1
                ## 1 0 1
                ## 0 1 1
                ## 1 1 1
                tstring = bin(i) ## this string has format '0bxxxxx'
                # print("tstring",tstring)
                tstring = tstring[2:] ## strip off leading "0b"
                tstring = "0"*(tlen-len(tstring))+tstring
                # print("tstring",tstring)
                tstring = tstring[1:] ## don't want to reverse first digit
                rstring = tstring[::-1]
                ## reverse the slice, this gives us the binary number we want
                s1 = "0"+rstring
                # print("s1",s1)
                ind = int(s1,2) ## where in the list we want to put this prob
                # print("ind",ind)
                # print("entry",entry)
                # print("1-entry",1-entry)
                # print("ind+plen",ind+plen)
                pset[ind] = entry
                pset[ind+plen] = 1-entry
            
#            for i in range(len(ptable)*2):
#                print "["+str(i)+"]", pset[i]; 
            return pset
        
        daifactors = dai.VecFactor() #dai.VectorFactor
        self.aNodes = toposort(self.ndict)
        self.node2label = dict()
        self.node2daiVar = dict()

        for graphid,node in enumerate(self.aNodes):
            self.node2label[node] = graphid
            self.node2daiVar[node] = dai.Var(graphid, node.nValues)
        
        for node,context in self.ndict.iteritems():
#            print "";
            vlist = []
            parents = context[0]
            ptable = context[1]
            if len(parents)==0:
                vset = self.node2daiVar[node]
#                print "Node:", node.name
            else:
                for parent in parents:
                    vlist.append(self.node2daiVar[parent])
#                    print "Parent:", parent.name;
                vlist.append(self.node2daiVar[node])
                vset = dai.VarSet(vlist[0],vlist[1])
                for node in vlist[2:]:
                    vset.append(node) #vset.add(node)
            f = dai.Factor(vset)
            daifactors.append(p2cpt(f, parents, ptable, self.node2label))        
        self.dainetwork = dai.FactorGraph(daifactors)
        self.dainetwork.WriteToFile("tempdai.fg")
        
    def setup2(self):
        def toposort(nodedict):
            visited = dict()
            nodelist = []
            def visit(node):
#                print node.name;
                try:
                    visited[node]
#                    print "visited!";
                except KeyError:
                    visited[node] = 1
                    if not(nodedict.has_key(node)): print "Error trying to visit:", node.name;
                    pset = nodedict[node][0] ## parents
                    for p in pset:
                        visit(p)
#                    print "Appending:", node.name;
                    nodelist.append(node)
            for node in nodedict.keys():
#                print "Starting with:", node.name;
                visit(node)
#            print "";
#            for node in nodelist:
#                print "Ending with:", node.name;
            return(nodelist)
        def p2cpt(pset, parents, ptable, node2label):
            def getIndex(aVals, aBase, aIndex):
                nIndex = 0;
                for i, nV in enumerate(aVals):
                    nF, nBases = 1, len(aBase[i+1:]);
                    for j in range(nBases):
                        nF *= aBase[aIndex[j+1]];
                    nIndex += (nV*nF);
                return nIndex;            
            def getPermutation(nIndex, aBase):
                aVals = [];
                for nB in aBase[::-1]:
                    aVals.append(nIndex % nB);
                    nIndex /= nB;
                aVals.reverse();
                return aVals;
            def permute(self,ptable, ids):
                order = [i[0] for i in sorted(enumerate(ids), key=lambda x:x[1])]
                pnew = [0]*len(ptable)
                for (i,val) in enumerate(ptable):    
                    d = self.tobin(i, len(ids))
                    od = [d[j] for j in order]
                    newi = self.frombin(od)
                    if newi>len(pnew)-1:
#                        newi = len(pnew)-1;
                        print "ERROR (newi>len(pnew)-1): Node's CPT has incorrect number of probability arguments.";
                    pnew[newi] = val
                return pnew            
            aV = [2,2,2];
            nVals = 1;
            for n in aV:
                nVals *= n;
#             aFrom, aTo = [0,1,2], [2,0,1];
#             aDiff = [aTo.index(i) for i in range(len(aTo))]; #[1,2,0]
#             aV2 = [aV[i] for i in aTo]; #[2,3,2]

            ids, aTo = [], [];
            for parent in parents:
                ids.append(node2label[parent])
                aTo.append(node2label[parent])
            ptable = permute(self, ptable,ids)
#            print "aTo:", aTo;

            plen = len(ptable)
#             tset = [0]*2*plen
            tlen = np.log2(plen)+1 ## length our binary string
            for i,entry in enumerate(ptable):
                entry = 1-entry
                ## entries all assume entry conditioned on others is 0
                ## ptables count up across the remaining variables in binary
                ## so e.g. for 3 variables (c|a,b), the table looks like
                ## 0 0 0
                ## 0 0 1
                ## 0 1 0
                ## 0 1 1
                ## however, dai requires that the full table be specified, and
                ## counts in "reversed" binary, so the full table for the above
                ## example would be ordered:
                ## 0 0 0
                ## 1 0 0
                ## 0 1 0
                ## 1 1 0
                ## 0 0 1
                ## 1 0 1
                ## 0 1 1
                ## 1 1 1
                tstring = bin(i) ## this string has format '0bxxxxx'
                # print("tstring",tstring)
                tstring = tstring[2:] ## strip off leading "0b"
                tstring = "0"*(tlen-len(tstring))+tstring
                # print("tstring",tstring)
                tstring = tstring[1:] ## don't want to reverse first digit
                rstring = tstring[::-1]
                ## reverse the slice, this gives us the binary number we want
                s1 = "0"+rstring
                # print("s1",s1)
                ind = int(s1,2) ## where in the list we want to put this prob
                # print("ind",ind)
                # print("entry",entry)
                # print("1-entry",1-entry)
                # print("ind+plen",ind+plen)
                pset[ind] = entry
                pset[ind+plen] = 1-entry
            
#            for i in range(len(ptable)*2):
#                print "["+str(i)+"]", pset[i]; 
            return pset
        
        self.aNodes = toposort(self.ndict)
        self.node2label = dict()
        for graphid,node in enumerate(self.aNodes):
            self.node2label[node] = graphid

    ##Run BP clamped and unclamped
    def postprob(self, nodes):
        p = []
        #Run Post (Clamped) with BP
        self.bp_clamped = dwrap.BPwrapper()
        for obs in self.obs:
            nid = self.node2label[obs[0]]
            self.bp_clamped.clamp(nid, obs[1])
        self.bp_clamped.run() 
        #Get Probs
        for node in nodes:
            self.bp_clamped.addVar(self.node2label[node])
            p.append(self.bp_clamped.getJoint(0))
        self.bp_clamped.resetVars()
        self.bp_clamped.dealloc();
        return p

    def priorprob(self, nodes):
        p = [];
        #Run Prior (Unclamped) with BP
        self.bp_unclamped = dwrap.BPwrapper();
        self.bp_unclamped.run();
        #Get Probs
        for node in nodes:
            self.bp_unclamped.addVar(self.node2label[node]);
            p.append(self.bp_unclamped.getJoint(0));
        self.bp_unclamped.resetVars();
        self.bp_unclamped.dealloc();        
        return p
    
    def getMAP(self):
        aMAP_States = [];
        #Run MAP (Clamped) with BP
        self.bp_clamped = dwrap.BPwrapper();
        for obs in self.obs:
            nid = self.node2label[obs[0]];
            self.bp_clamped.clamp(nid, obs[1]);
        self.bp_clamped.run();
        nLogZ = self.bp_clamped.getLogZ();
        print "LogZ:", nLogZ, "Z:", math.exp(nLogZ);
        self.bp_clamped.runMAP();
        nLogScore = self.bp_clamped.getLogScoreMAP();
        print "LogScore:", nLogScore, "Score:", math.exp(nLogScore);
#         print "P(MAP|E):", math.exp(nLogScore)/math.exp(nLogZ);
        #Get MAP
        for i in range(self.bp_clamped.getStateMAP_Size()):
            aMAP_States.append(self.bp_clamped.getStateMAP(i));
        self.bp_clamped.resetVars();
        self.bp_clamped.dealloc();
        return aMAP_States;

    def write2dot(self,fname = "graph.dot", plabels=False):
        node2label = self.node2label
        f = open(fname, "w")
        f.write("digraph G {\n")
        #f.write("node[shape=circle, width=0.4,fixedsize=true];\n")
        f.write("node[shape=circle, width=0.4];\n")
        for node in self.ndict.keys():
            if plabels:
                l = ""
                for i in self.ndict[node][1]:                
                    l += "x"+str(Decimal(str(10*i)).quantize(Decimal('1'), rounding = ROUND_DOWN))
            else:
                l = "\""+ node.name +"\""
            f.write(str(node2label[node]))
            if node in map(lambda x : x[0], self.obs):
                f.write("[label="+l+",style=filled,color=blue]")
#             elif node.name == NN.WAITING:
#                 f.write("[label="+l+",style=filled,color=green]")
#             elif node.name == NN.STAY_PUT:
#                 f.write("[label="+l+",style=filled,color=yellow]")
#             elif node.name == NN.HOLDING_OBJ:
#                 f.write("[label="+l+",style=filled,color=black]")
#             elif node.name == NN.EXCHANGE_DROP:
#                 f.write("[label="+l+",style=filled,color=red]")
            else:
                #f.write("[label="+self.name(node)+"]")
                f.write("[label="+l+"]")
            f.write(";\n")
            for parent in self.ndict[node][0]:
                f.write(str(node2label[parent])+" -> "+ str(node2label[node])+";\n")
                #f.write(self.name(parent)+" -> "+ self.name(node)+";\n")        
        f.write("}")
        f.close()

    def tobin(self,entry, n):
        digits = []
        ndigits = 0
        while (entry > 0):
            digits = [entry - 2*int(float(entry)/float(2))]+digits
            ndigits = ndigits+1
            entry = entry/2        
        while (ndigits < n):
            digits = [0]+digits
            ndigits = ndigits+1
        return(digits)

    def frombin(self,digits):
        entry = 0
        digits.reverse()
        for i,d in enumerate(digits):
            entry = pow(2,i)*d +entry
        return(entry)

    def name(self, nodes):
        def nname(node):
            if node.extendable:
                sstr = ""
                for s in node.s:
                    sstr = sstr+('%03d'%s)
                ostr = ""
                for o in node.o:
                    ostr = ostr+('%03d'%o)
                tstr = ""
                for t in node.t:
                    tstr = tstr+('%03d'%t)
                return(node.name+"_"+sstr+"_"+ostr+"_"+tstr)
            else:
                return(node.name)
            
        if type(nodes).__name__=='list':
            return(map(nname, nodes))
        return(nname(nodes))
